import getpass
import os
from langchain.chat_models import init_chat_model
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_community.embeddings import DashScopeEmbeddings
from dotenv import load_dotenv


load_dotenv(".env.local")

# if not os.environ.get("OPENAI_API_KEY"):
#     os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter API key for OpenAI: ")

# model = init_chat_model("gpt-4o-mini", model_provider="openai")

file_path = "/Users/leisurexi/workspace/python/langchain-doc-qa/example_data/nke-10k-2023.pdf"
loader = PyPDFLoader(file_path)
docs = loader.load()
print(len(docs))

# 文本拆分器
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000, chunk_overlap=200, add_start_index=True
)
all_splits = text_splitter.split_documents(docs)

# 嵌入
embeddings = DashScopeEmbeddings(model="text-embedding-v1")

# 向量存储
vector_store = InMemoryVectorStore(embeddings)
ids = vector_store.add_documents(documents=all_splits)

# 根据与字符串查询的相似性返回文档
results = vector_store.similarity_search(
    "How many distribution centers does Nike have in the US?"
)
# 异步查询
#results = await vector_store.asimilarity_search("When was Nike incorporated?")
print(results[0])
